{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Notebook for preparing and saving SBM_PATTERN graphs in DGL form"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "import pickle\n",
    "import time\n",
    "import os\n",
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Download SBM_PATTERN dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.path.isfile('SBM_PATTERN.zip'):\n",
    "    print('downloading..')\n",
    "    !curl https://www.dropbox.com/s/qvu0r11tjyt6jyb/SBM_PATTERN.zip?dl=1 -o SBM_PATTERN.zip -J -L -k\n",
    "    !unzip SBM_PATTERN.zip -d ./\n",
    "    !rm -r __MACOSX/\n",
    "else:\n",
    "    print('File already downloaded')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Convert to DGL format and save with pickle"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/home/vijay/graphdeeplearning/benchmarking-gnns\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "os.chdir('../../') # go to root folder of the project\n",
    "print(os.getcwd())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using backend: pytorch\n"
     ]
    }
   ],
   "source": [
    "import pickle\n",
    "\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "from data.SBMs import SBMsDatasetDGL \n",
    "\n",
    "from data.data import LoadData\n",
    "from torch.utils.data import DataLoader\n",
    "from data.SBMs import SBMsDataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DotDict(dict):\n",
    "    def __init__(self, **kwds):\n",
    "        self.update(kwds)\n",
    "        self.__dict__ = self"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "start = time.time()\n",
    "\n",
    "DATASET_NAME = 'SBM_PATTERN'\n",
    "dataset = SBMsDatasetDGL(DATASET_NAME) \n",
    "\n",
    "print('Time (sec):',time.time() - start) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "10000\n",
      "2000\n",
      "2000\n",
      "(Graph(num_nodes=108, num_edges=4884,\n",
      "      ndata_schemes={'feat': Scheme(shape=(), dtype=torch.int64)}\n",
      "      edata_schemes={'feat': Scheme(shape=(1,), dtype=torch.float32)}), tensor([0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0,\n",
      "        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0,\n",
      "        0, 1, 0, 0, 0, 0, 0, 1, 0, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1,\n",
      "        0, 0, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0,\n",
      "        1, 0, 0, 0, 1, 0, 1, 1, 0, 0, 0, 0], dtype=torch.int16))\n",
      "(Graph(num_nodes=108, num_edges=4738,\n",
      "      ndata_schemes={'feat': Scheme(shape=(), dtype=torch.int64)}\n",
      "      edata_schemes={'feat': Scheme(shape=(1,), dtype=torch.float32)}), tensor([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0,\n",
      "        0, 1, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
      "        1, 0, 1, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 1,\n",
      "        0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0,\n",
      "        0, 0, 0, 0, 0, 1, 0, 1, 0, 0, 1, 1], dtype=torch.int16))\n",
      "(Graph(num_nodes=94, num_edges=3772,\n",
      "      ndata_schemes={'feat': Scheme(shape=(), dtype=torch.int64)}\n",
      "      edata_schemes={'feat': Scheme(shape=(1,), dtype=torch.float32)}), tensor([0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 1,\n",
      "        1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
      "        1, 0, 1, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 1,\n",
      "        0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 0, 1, 0, 0],\n",
      "       dtype=torch.int16))\n"
     ]
    }
   ],
   "source": [
    "print(len(dataset.train))\n",
    "print(len(dataset.val))\n",
    "print(len(dataset.test))\n",
    "\n",
    "print(dataset.train[0])\n",
    "print(dataset.val[0])\n",
    "print(dataset.test[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "start = time.time()\n",
    "\n",
    "with open('data/SBMs/SBM_PATTERN.pkl','wb') as f:\n",
    "        pickle.dump([dataset.train,dataset.val,dataset.test],f)\n",
    "        \n",
    "print('Time (sec):',time.time() - start) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[I] Loading dataset SBM_PATTERN...\n",
      "train, test, val sizes : 10000 2000 2000\n",
      "[I] Finished loading.\n",
      "[I] Data load time: 12.3468s\n"
     ]
    }
   ],
   "source": [
    "DATASET_NAME = 'SBM_PATTERN'\n",
    "dataset = LoadData(DATASET_NAME) \n",
    "trainset, valset, testset = dataset.train, dataset.val, dataset.test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Time (sec): 0.0004932880401611328\n"
     ]
    }
   ],
   "source": [
    "start = time.time()\n",
    "\n",
    "batch_size = 10\n",
    "collate = SBMsDataset.collate\n",
    "train_loader = DataLoader(trainset, batch_size=batch_size, shuffle=True, collate_fn=collate)\n",
    "\n",
    "print('Time (sec):',time.time() - start) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
